#include <vector>
#include <cmath>
#include <algorithm>
#include <mutex>

// Define precision type (double precision in original)
using Real = double;
constexpr Real ZERO = 0.0;
constexpr Real ONE = 1.0;
constexpr Real ZERO3 = 0.3;

class ContactForceCalculator {
public:
    // Main force calculation routine
    void calculateForces(
        const std::vector<std::array<Real, 3>>& nodeCoords,    // Node coordinates (X)
        std::vector<Real>& forces,                            // Force array (E)
        const std::vector<std::array<int, 6>>& segmentConnect, // Segment connectivity (IRECT)
        const std::vector<int>& masterNodes,                  // MSR array
        const std::vector<int>& slaveNodeList,                // NSV array
        const std::vector<int>& nodeLocations,                // ILOC array
        const std::vector<int>& contactSegments,              // IRTL array
        const std::vector<Real>& stiffness,                   // STF array
        std::vector<std::array<Real, 3>>& frictionForces,     // FROLD array
        Real contactThickness,                                // THKSV
        Real slaveDamping,                                   // SDMPS
        Real masterDamping,                                  // SDMPM
        Real drawbeadStress,                                // SIGDB
        Real normalStress,                                  // SIGND
        int drawbeadFlag,                                   // IDRAW
        Real& slaveAreaSum,                                // ARSL
        const std::vector<Real>& sectionAreas,            // ASEC array
        const std::vector<Real>& velocities,             // V array
        const std::vector<Real>& nodalMasses,           // XMS array
        std::vector<Real>& contactFriction,            // CTFR array
        const std::vector<Real>& nodalThickness,       // THKNO array
        const std::vector<std::array<Real, 3>>& nodalNormals, // XNORM array
        const std::vector<Real>& slaveAreas,           // ARST array
        int activationFlag,                          // IACT
        int normalFlag,                             // INOR
        Real frictionCoeff,                        // FRIC
        const std::vector<std::array<Real, 3>>& bendDirections, // XBEND array
        Real& springWork,                        // SPRWRK
        Real& frictionWork,                     // DFRW
        Real timeStep,                         // DTI
        const std::vector<Real>& bendFactors, // CBEND array
        std::vector<Real>& tangentFriction,  // CTANPR array
        int startIndex,                     // IBEG
        int endIndex                       // IEND
    ) {
        if (drawbeadFlag == 3) return;

        constexpr int LGR = 32; // Group size for vectorization
        int numPairs = endIndex - startIndex + 1;
        int numGroups = (numPairs + LGR - 1) / LGR;
        int groupSize = LGR;

        // Thread safety mechanism
        std::mutex forceUpdateMutex;

        for (int group = 0; group < numGroups; ++group) {
            int groupStart = startIndex + group * LGR;
            if (group == numGroups - 1) {
                groupSize = numPairs - LGR * (numGroups - 1);
            }

            // Process each contact pair in the group
            processContactGroup(
                nodeCoords, forces, segmentConnect, masterNodes, slaveNodeList,
                nodeLocations, contactSegments, stiffness, frictionForces,
                contactThickness, slaveDamping, masterDamping, drawbeadStress,
                normalStress, drawbeadFlag, slaveAreaSum, sectionAreas,
                velocities, nodalMasses, contactFriction, nodalThickness,
                nodalNormals, slaveAreas, activationFlag, normalFlag,
                frictionCoeff, bendDirections, springWork, frictionWork,
                timeStep, bendFactors, tangentFriction, groupStart, groupSize,
                forceUpdateMutex);
        }
    }

private:
    // Helper methods
    
    void processContactGroup(
        const std::vector<std::array<Real, 3>>& nodeCoords,
        std::vector<Real>& forces,
        const std::vector<std::array<int, 6>>& segmentConnect,
        const std::vector<int>& masterNodes,
        const std::vector<int>& slaveNodeList,
        const std::vector<int>& nodeLocations,
        const std::vector<int>& contactSegments,
        const std::vector<Real>& stiffness,
        std::vector<std::array<Real, 3>>& frictionForces,
        Real contactThickness,
        Real slaveDamping,
        Real masterDamping,
        Real drawbeadStress,
        Real normalStress,
        int drawbeadFlag,
        Real& slaveAreaSum,
        const std::vector<Real>& sectionAreas,
        const std::vector<Real>& velocities,
        const std::vector<Real>& nodalMasses,
        std::vector<Real>& contactFriction,
        const std::vector<Real>& nodalThickness,
        const std::vector<std::array<Real, 3>>& nodalNormals,
        const std::vector<Real>& slaveAreas,
        int activationFlag,
        int normalFlag,
        Real frictionCoeff,
        const std::vector<std::array<Real, 3>>& bendDirections,
        Real& springWork,
        Real& frictionWork,
        Real timeStep,
        const std::vector<Real>& bendFactors,
        std::vector<Real>& tangentFriction,
        int groupStart,
        int groupSize,
        std::mutex& forceUpdateMutex
    ) {
        // Local storage for group processing
        std::vector<int> activeIndices;
        std::vector<std::array<std::array<Real, 3>, 4>> segmentNodes(groupSize);
        std::vector<std::array<Real, 3>> slavePositions(groupSize);
        std::vector<Real> thickness(groupSize);
        std::vector<Real> penetration(groupSize);
        std::vector<std::array<Real, 3>> normals(groupSize);
        std::vector<std::array<std::array<Real, 3>, 4>> nodeVelocities(groupSize);
        std::vector<std::array<Real, 3>> relativeVelocities(groupSize);
        std::vector<Real> normalForces(groupSize);
        std::vector<std::array<Real, 3>> contactForces(groupSize);

        // First pass: collect data and check for active contacts
        for (int kk = 0; kk < groupSize; ++kk) {
            int ii = groupStart + kk;
            int slaveNode = slaveNodeList[ii];
            
            // Skip invalid nodes
            if (slaveNode <= 0) continue;

            int segmentId = std::abs(contactSegments[ii]);
            
            // Skip if no contact segment
            if (segmentId == 0) {
                frictionForces[ii] = {0.0, 0.0, 0.0};
                continue;
            }

            // Store segment node data
            for (int j = 0; j < 4; ++j) {
                int nodeId = segmentConnect[segmentId][j];
                segmentNodes[kk][j] = nodeCoords[nodeId];
            }

            // Store slave node position
            slavePositions[kk] = nodeCoords[slaveNode];

            // Determine thickness
            if (activationFlag == 1) {
                thickness[kk] = 0.5 * nodalThickness[slaveNode];
            } else {
                thickness[kk] = contactThickness;
            }

            activeIndices.push_back(kk);
        }

        // Calculate penetration and normals for active contacts
        calculatePenetrationAndNormals(
            segmentNodes, slavePositions, thickness, penetration, normals,
            activeIndices);

        // Process active contacts
        for (int kk : activeIndices) {
            int ii = groupStart + kk;
            int slaveNode = slaveNodeList[ii];

            // Skip if no penetration (except for drawbead cases)
            if (penetration[kk] >= 0.0) {
                if (drawbeadFlag == 0) {
                    frictionForces[ii] = {0.0, 0.0, 0.0};
                    continue;
                } else {
                    penetration[kk] = -1.0e-7;
                }
            }

            // Limit penetration to reasonable values
            Real penetrationLimit;
            if (thickness[kk] == 0.0) {
                Real area = calculateSegmentArea(segmentNodes[kk]);
                penetrationLimit = -0.20 * std::sqrt(0.125 * area);
            } else {
                penetrationLimit = -10.0 * thickness[kk];
            }

            if (penetration[kk] < penetrationLimit) {
                frictionForces[ii] = {0.0, 0.0, 0.0};
                continue;
            }

            penetration[kk] = std::max(0.1 * penetrationLimit, penetration[kk]);
        }

        // Calculate velocities and forces
        calculateVelocitiesAndForces(
            segmentNodes, slavePositions, velocities, nodalMasses, stiffness,
            penetration, normals, slaveDamping, masterDamping, drawbeadFlag,
            bendDirections, normalForces, contactForces, relativeVelocities,
            nodeVelocities, activeIndices, groupStart, slaveNodeList);

        // Update forces with thread safety
        {
            std::lock_guard<std::mutex> lock(forceUpdateMutex);
            updateForceContributions(
                forces, contactFriction, sectionAreas, slaveAreas,
                normalForces, contactForces, penetration, springWork,
                activeIndices, groupStart, slaveNodeList, segmentNodes,
                nodeVelocities, masterDamping);
        }

        // Handle friction calculations if enabled
        if (frictionCoeff != 0.0 || drawbeadFlag != 0) {
            calculateFrictionForces(
                frictionForces, relativeVelocities, stiffness, timeStep,
                frictionCoeff, normalForces, slaveAreas, drawbeadFlag,
                bendDirections, bendFactors, drawbeadStress, normalStress,
                forces, sectionAreas, frictionWork, activeIndices, groupStart,
                slaveNodeList, segmentNodes, nodeVelocities, forceUpdateMutex);
        }
    }


};